/*
 * FlvMediaMuxer.cpp
 *
 *  Created on: 2020年6月6日
 *      Author: chuanjiang.zh
 */

#include "FlvMediaMuxer.h"

#include "flv-writer.h"
#include "flv-muxer.h"
#include "flv-proto.h"

#include "CLog.h"
#include "AacHelper.h"


class FlvMediaMuxerPrivate
{
public:
	FlvMediaMuxerPrivate():
		m_muxer(),
		m_writer()
	{
	}

	void firePacket(UMediaPacket& pkt)
	{
		if (m_sink)
		{
			m_sink(pkt);
		}
	}

	void setSink(UMediaMuxer::Sink sink)
	{
		m_sink = sink;
	}


	UMediaMuxer::Sink	m_sink;

	flv_muxer_t*	m_muxer;
	void*	m_writer;

};



FlvMediaMuxer::FlvMediaMuxer() :
	m_pimp(new FlvMediaMuxerPrivate())
{
}

FlvMediaMuxer::~FlvMediaMuxer()
{
	close();

	delete m_pimp;
}

bool FlvMediaMuxer::open(const UMediaFormat& fmt)
{
    if (!fmt.isValid())
    {
        return false;
    }

    int audioCodecId = getAudioCodecId(fmt.m_audioCodec);
    int videoCodecId = getVideoCodecId(fmt.m_codec);

    if (videoCodecId <= 0)
    {
        return false;
    }

    m_format = fmt;

	m_pimp->m_muxer = flv_muxer_create(flv_onmuxer, this);
	m_pimp->m_writer = flv_writer_create2(flv_write, this);

	struct flv_metadata_t metadata = flv_metadata_t();
    
    if (audioCodecId > 0)
    {
        metadata.audiocodecid = audioCodecId;
        metadata.audiodatarate = fmt.m_audioBitrate;

        //if (fmt.m_audioCodec == UMEDIA_CODEC_AAC)
        metadata.audiosamplerate = fmt.m_sampleRate;
        metadata.audiosamplesize = 16;
        metadata.stereo = (fmt.m_channels > 1);
    }

	metadata.videocodecid = getVideoCodecId(fmt.m_codec);
	metadata.videodatarate = fmt.m_bitrate;
	metadata.framerate = fmt.m_framerate;
	metadata.width = fmt.m_width;
	metadata.height = fmt.m_height;

	int ret = flv_muxer_metadata(m_pimp->m_muxer, &metadata);

	return (ret == 0);
}

void FlvMediaMuxer::close()
{
    if (m_pimp->m_writer)
    {
        flv_writer_destroy(m_pimp->m_writer);
        m_pimp->m_writer = nullptr;
    }

    if (m_pimp->m_muxer)
    {
        flv_muxer_destroy(m_pimp->m_muxer);
        m_pimp->m_muxer = nullptr;
    }

}

bool FlvMediaMuxer::isOpen()
{
	return m_pimp->m_muxer != nullptr;
}


bool FlvMediaMuxer::getOutFormat(UMediaFormat& fmt)
{
	fmt = m_format;
	return m_format.isValid();
}

bool FlvMediaMuxer::input(UMediaPacket& pkt)
{
	if (pkt.size() <= 0)
	{
		return false;
	}

	bool done = false;
	if (pkt.isVideo())
	{
		int64_t duration = 0;
		int64_t ts = m_videoShift.shift(pkt.pts(), duration);
		uint32_t pts = (uint32_t)(ts / 1000);

        //CLog::debug("FlvMediaMuxer::input. pts: %lld, new pts: %lld, ts: %lld\n", pkt.pts(), ts, pts);

        if (m_format.m_codec == UMEDIA_CODEC_H264)
        {
            flv_muxer_avc(m_pimp->m_muxer, pkt.data(), pkt.size(), pts, pts);
        }
        else if (m_format.m_codec == UMEDIA_CODEC_HEVC)
        {
            flv_muxer_hevc(m_pimp->m_muxer, pkt.data(), pkt.size(), pts, pts);
        }
        else
        {
            //
        }

		done = true;
	}
	else if (pkt.isAudio())
	{
		int64_t duration = 0; 
        int64_t ts = m_audioShift.shift(pkt.pts(), duration);
        uint32_t pts = (uint32_t)(ts / 1000);

        if (m_format.m_audioCodec == UMEDIA_CODEC_AAC)
        {
            if (AacHelper::isAdts(pkt.data(), pkt.size()))
            {
                flv_muxer_aac(m_pimp->m_muxer, pkt.data(), pkt.size(), pts, pts);
            }
            else
            {
                int pktSize = pkt.size() + AacHelper::ADTS_HEADER_SIZE;
                UMediaPacket adtsPacket;
                adtsPacket.newPacket(pktSize);
				memcpy(adtsPacket.data() + AacHelper::ADTS_HEADER_SIZE, pkt.data(), pkt.size());

                AacHelper::makeAdtsHeader(m_format.m_audioProfile, m_format.m_sampleRate, m_format.m_channels,
                    pktSize, adtsPacket.data(), adtsPacket.size());
                
                flv_muxer_aac(m_pimp->m_muxer, adtsPacket.data(), adtsPacket.size(), pts, pts);
            }
        }
        else if (m_format.m_audioCodec == UMEDIA_CODEC_MP3)
        {
            flv_muxer_mp3(m_pimp->m_muxer, pkt.data(), pkt.size(), pts, pts);
        }
        else
        {
            //
        }

		done = true;
	}
	else
	{
		//
	}
	return done;
}

void FlvMediaMuxer::setSink(Sink sink)
{
	m_pimp->setSink(sink);
}

int FlvMediaMuxer::flv_onmuxer(void* flv, int type, const void* data, size_t bytes, uint32_t timestamp)
{
	FlvMediaMuxer* pthis = (FlvMediaMuxer*)flv;
	pthis->onmuxer(type, data, bytes, timestamp);

	return 0;
}

void FlvMediaMuxer::onmuxer(int type, const void* data, size_t bytes, uint32_t timestamp)
{
	flv_writer_input(m_pimp->m_writer, type, data, bytes, timestamp);
}

int FlvMediaMuxer::getVideoCodecId(int codec)
{
    if (codec == UMEDIA_CODEC_H264)
    {
        return FLV_VIDEO_H264;
    }
    else if (codec == UMEDIA_CODEC_HEVC)
    {
        return FLV_VIDEO_H265;
    }

    return 0;
}

int FlvMediaMuxer::getAudioCodecId(int codec)
{
    if (codec == UMEDIA_CODEC_G711A)
    {
        return FLV_AUDIO_G711A;
    }
    else if (codec == UMEDIA_CODEC_G711U)
    {
        return FLV_AUDIO_G711U;
    }
    else if (codec == UMEDIA_CODEC_MP3)
    {
        return FLV_AUDIO_MP3;
    }
    else if (codec == UMEDIA_CODEC_AAC)
    {
        return FLV_AUDIO_AAC;
    }

	return 0;
}

int FlvMediaMuxer::flv_write(void* param, const void* buf, int len)
{
	FlvMediaMuxer* pthis = (FlvMediaMuxer*)param;
	return pthis->onwrite(buf, len);
}

int FlvMediaMuxer::onwrite(const void* buf, int len)
{
	UMediaPacket packet;
	packet.newPacket(len);
	packet.set_type(UMEDIA_TYPE_DATA);

	int64_t ts = 0;
	packet.set_pts(ts);
	packet.set_dts(ts);

	packet.copy((const uint8_t*)buf, len);

	m_pimp->firePacket(packet);

	return len;
}

const char* FlvMediaMuxer::getName()
{
    return "flv";
}